"""
@Time: 2021/1/25 下午 4:56
@Author: jinzhuan
@File: ner_ace2005_toolkit.py
@Desc: 
"""

import torch
from cognlp import *
torch.cuda.set_device(5)

tokenize_toolkit = TokenizeToolkit()
words = tokenize_toolkit.run("It's irrelevant because the REAL question is whether an active board member can, while on the board, arrange to be given money by the USCF.")

ner_toolkit = NerAce2005Toolkit(
    bert_model='bert-base-cased',
    model_path='../../../cognlp/data/ner/ace2005/toolkit/model.pt',
    vocabulary_path='../../../cognlp/data/ner/ace2005/toolkit/vocabulary.txt',
    device=torch.device("cuda"),
    device_ids=[5],
    max_seq_length=256)
ner_result = ner_toolkit.run(words)
print(ner_result)
